from langdev_helper.llm.lcex import llm_lcex as llm


from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from typing import List, Sequence
from langgraph.graph import END, MessageGraph

from chain import generate_chain, reflect_chain


REFLECT = 'reflect'
GENERATE = 'generate'

def generation_node(state: Sequence[BaseMessage]):
    return generate_chain.invoke({"messages": state})

def reflection_node(messages: Sequence[BaseMessage]) -> List[BaseMessage]:
    cls_map = {"ai": HumanMessage, "human": AIMessage}
    translated = [messages[0]] + [
        cls_map[msg.type](content=msg.content) for msg in messages[1:]
    ]
    res = reflect_chain.invoke({"messages": translated})
    return HumanMessage(content=res.content)

builder = MessageGraph()
builder.add_node(GENERATE, generation_node)
builder.add_node(REFLECT, reflection_node)
builder.set_entry_point(GENERATE)

def should_continue(state: List[BaseMessage]):
    if len(state) > 3:
        return END
    return REFLECT

builder.add_conditional_edges(GENERATE, should_continue)
builder.add_edge(REFLECT, GENERATE)
graph = builder.compile()
graph.get_graph().print_ascii()

print(graph.get_graph().draw_mermaid())

import pprint
pp = pprint.PrettyPrinter(indent=4)


n=1
for event in graph.stream(
    [
        HumanMessage(
            content='就《小王子》为何与现代童年息息相关撰写一篇文章'
        )
    ],
):
    pp.pprint(event)
    print(f"---: {n}")
    n= n+1
